# 研究信噪比与距离间关系任务
import json
import numpy as np
from typing import List, Dict, Any
import matplotlib.pyplot as plt
from apps.cpm.ce.atom_ce import AtomCe
from apps.cpm.cpm_engine import CpmEngine
from apps.cpm.tes.range_utils import RangeUtils
from apps.cpm.tes.snr_range_utils import SnrRangeUtils
from core.radar_core import RadarCore
from common.app_registry import AppRegistry as AC

class SnrRangeAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.tes.tes_snr_range.TesSnrRange'

    def startup(self, query:str) -> Dict:
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        atom_prompt = AtomCe.generate(query=query)
        print(f'{atom_prompt}')
        resp = CpmEngine.infer(query=atom_prompt)
        params = json.loads(resp[8:-3])
        # 规划执行步骤
        plans = self.plan_processing(params)
        for plan in plans:
            print(f'plan: {plan};')
        # fig, ax, axs = SnrRangeUtils.init_matplotlib(rows=2, cols=1)
        # generate_range 功能调用
        range_values = RangeUtils.divide_ranges(params['ranges'][0], params['ranges'][1], params['ranges'][2])
        # snr_range_rcs
        pt = params['pt']
        freq = params['freq']
        g = params['g']
        te = params['te']
        b = params['b']
        nf = params['nf']
        loss = params['loss']
        ranges_km = range_values/1000.0
        sigma_values = params['sigma']
        # 内部Agent调用
        snr_matrix1 = np.array([RadarCore.radar_eq(pt[0], freq, g, s, te, b, nf, loss, ranges_km) 
                            for s in sigma_values])
        # MCP服务器调用
        SnrRangeUtils.create_radar_plot(AC.axs[0], ranges_km, snr_matrix1, title='信号噪比距离关系（不同RCS）', labels=[r'$\sigma$ = 0 dBsm', 
                r'$\sigma$ = -10 dBsm', 
                r'$\sigma$ = +10 dBsm'])
        # snr_range_pt
        snr_matrix2 = np.array([RadarCore.radar_eq(p, freq, g, sigma_values[0], te, b, nf, loss, ranges_km) 
                            for p in pt])
        # 绘制第二个图（不同功率）
        SnrRangeUtils.create_radar_plot(
            AC.axs[1], ranges_km, snr_matrix2,
            title='信噪比与距离关系图（不同发射功率）',
            labels=['Pt = 1.5 MW', 
                'Pt = 0.6 MW', 
                'Pt = 2.7 MW']
        )
        # 统一调整
        plt.tight_layout()
        plt.show()
        return params

    def run(self, input:Dict) -> Dict:
        pass

    def plan_processing(self, params:Any) -> List[Dict]:
        print(f'######### params: {type(params)}; ???????????')
        plans = [
            {'desc': '生成距离列表值', 'params': {'start': params['ranges'][0], 'end': params['ranges'][1], 'step': params['ranges'][2]}},
            {'desc': '计算snr range rcs', 'params': params},
            {'desc': '计算snr range pt', 'params': params},
            {'desc': '绘制曲线', 'params': None}
        ]
        return plans